/*!
  @file dense_stereo_graphcut.h
  @copyright 2013 Kubota Lab. All rights resereved.
*/

#ifndef _DENSE_STEREO_GRAPHCUT_H_
#define _DENSE_STEREO_GRAPHCUT_H_

#include <ros/ros.h>

#include <image_geometry/stereo_camera_model.h>
#include <image_transport/image_transport.h>
#include <image_transport/subscriber_filter.h>
#include <message_filters/subscriber.h> 
#include <message_filters/synchronizer.h>
#include <message_filters/sync_policies/approximate_time.h>

#include <sensor_msgs/CameraInfo.h>
#include <sensor_msgs/Image.h>
#include <sensor_msgs/image_encodings.h>
#include <stereo_msgs/DisparityImage.h>
#include <std_msgs/Int8.h>
#include <std_msgs/Float32.h>

#include <boost/thread/mutex.hpp>

#include <opencv2/core/core.hpp>

#include "../maxflow-v3.01/graph.h"
//#include "maxflow-v3.01/instances.inc"

namespace stereo {

/*! 
 *  @brief  Dense Stereo
 *  @author Kyohei Otsu <kyon@ac.jaxa.jp>
 *  @date   2013
 *  
 *  Performs dense stereo 
 */

class DenseStereoGraphcut
{
 public:
  typedef message_filters::sync_policies::ApproximateTime<sensor_msgs::Image, sensor_msgs::CameraInfo, sensor_msgs::Image, sensor_msgs::CameraInfo> SyncPolicy;
  typedef Graph<double,double,double> GraphType;


  explicit DenseStereoGraphcut(ros::NodeHandle &nh, ros::NodeHandle &private_nh) 
      : nh_(nh)
      , private_nh_(private_nh)
      , it_(nh_)
      , processing_(false)
      , trigger_(false)
      , scale_(1.0)
      , wsmooth_(0.5)
      , winsize_(5)
      , distmax_(10.0)
      , distmin_(1.0)
      , dispmax_(-1)
      , dispmin_(-1)
      , do_sobel_(false)
  {
    // initializing pub/sub, server
    imgL_sub_.subscribe(it_, nh_.resolveName("left/image_rect"), 3);
    imgR_sub_.subscribe(it_, nh_.resolveName("right/image_rect"), 3);
    infoL_sub_.subscribe(nh_, nh_.resolveName("left/camera_info"), 3);
    infoR_sub_.subscribe(nh_, nh_.resolveName("right/camera_info"), 3);
    sync_.reset(new message_filters::Synchronizer<SyncPolicy>(SyncPolicy(10), imgL_sub_, infoL_sub_, imgR_sub_, infoR_sub_));
    sync_->registerCallback(boost::bind(&DenseStereoGraphcut::stereoImageCB, this, _1, _2, _3, _4));
    disp_pub_ = nh_.advertise<stereo_msgs::DisparityImage>(nh_.resolveName("disparity/image"), 3);

    trigger_sub_ = nh_.subscribe(nh_.resolveName("conf/trigger"), 3, &DenseStereoGraphcut::triggerCB, this);
    smooth_sub_ = nh_.subscribe(nh_.resolveName("conf/smoothness"), 3, &DenseStereoGraphcut::smoothnessCB, this);
    scale_sub_ = nh_.subscribe(nh_.resolveName("conf/image_scale"), 3, &DenseStereoGraphcut::scaleCB, this);


    // initializing optional parameters
    // e.g.) private_nh_.param<TYPE>("PARAM_NAME", PARAM_VAR, DEFAULT_VALUE);
    private_nh_.param<double>("scale", scale_, scale_);
    private_nh_.param<double>("smoothness", wsmooth_, wsmooth_);
    private_nh_.param<int>("winsize", winsize_, winsize_);
    private_nh_.param<bool>("do_sobel", do_sobel_, do_sobel_);

    private_nh_.param<double>("max_distance", distmax_, distmax_);
    private_nh_.param<double>("min_distance", distmin_, distmin_);
  }

  ~DenseStereoGraphcut() { }

  void stereoImageCB(
      const sensor_msgs::ImageConstPtr &imgL, const sensor_msgs::CameraInfoConstPtr &infoL, 
      const sensor_msgs::ImageConstPtr &imgR, const sensor_msgs::CameraInfoConstPtr &infoR);

  void triggerCB(const std_msgs::Int8ConstPtr &msg)
  {
    boost::mutex::scoped_lock lock(mutex_);
    trigger_ = true;
    ROS_INFO("DenseStereo: Received trigger");
  }

  void smoothnessCB(const std_msgs::Float32ConstPtr &msg)
  {
    ROS_INFO("smoothness");
    boost::mutex::scoped_lock lock(mutex_);
    while (processing_)
    {
      ros::Duration(1).sleep();
      ros::spinOnce();
    }
    wsmooth_ = msg->data;
    ROS_INFO_STREAM("DenseStereo: Smoothness is set to " << wsmooth_);
  }

  void scaleCB(const std_msgs::Float32ConstPtr &msg)
  {
    boost::mutex::scoped_lock lock(mutex_);
    while (processing_)
    {
      ros::Duration(1).sleep();
      ros::spinOnce();
    }
    scale_ = msg->data;
    ROS_INFO_STREAM("DenseStereo: Image scale is set to " << scale_);
  }



 private:
  void process(const cv::Mat &imgL, const cv::Mat &imgR, cv::Mat &data);

  void performSobel(const cv::Mat &src, cv::Mat &dst);

  //! smoothness term
  double wdiff(const float fp, const float fq, const float u=1.0)
  {
    //return (fabs(fp - fq) < 1.0)? 0: u;
    return 4 * std::min(9, abs(fp - fq));
  }
  double paramPotts(const uchar &Ip, const uchar &Iq)
  {
    return (abs(Ip - Iq) <= 5)? 2 * wsmooth_: wsmooth_;
  }

  //! data term
  double sad(const int x, const int y, const cv::Mat &imL, const cv::Mat &imR, const int dsp);
  double ad(const int x, const int y, const cv::Mat &imL, const cv::Mat &imR, float &d, bool rewrite=false);

  template <typename T, int N>
  cv::Vec<T, N> interpolateN(const cv::Mat &img, const double x, const double y)
  {
    if(x < 0 && x > img.cols - 1 || y < 0 && y > img.rows - 1)
    {
      const cv::Vec<uchar, N>& a = img.at<cv::Vec<uchar, N> >(floor(y), floor(x));
      cv::Vec<T, N> c;
      for (int i = 0; i < N; ++i) 
        c[i] = a[i];
      return c;
    }

    double eps = 0.001;
    double xf = floor(x);
    double yf = floor(y);
    double xc = (x - xf < eps)? xf: xf + 1;
    double yc = (y - yf < eps)? yf: yf + 1;

    const cv::Vec<uchar, N>& a = img.at<cv::Vec<uchar, N> >(yf, xf);
    const cv::Vec<uchar, N>& b = img.at<cv::Vec<uchar, N> >(yc, xc);
    cv::Vec<T, N> c;
    for (int i = 0; i < N; ++i) 
      c[i] = a[i] + (b[i] - a[i]) * (x - xf);
    return c;
  }

  double interpolate(const cv::Mat &img, const double x, const double y)
  {
    if(x < 0 && x > img.cols - 1 || y < 0 && y > img.rows - 1)
      return img.at<uchar>(floor(y), floor(x));

    double eps = 0.001;
    double xf = floor(x);
    double yf = floor(y);
    double xc = (x - xf < eps)? xf: xf + 1;
    double yc = (y - yf < eps)? yf: yf + 1;

    const uchar& a = img.at<uchar>(yf, xf);
    const uchar& b = img.at<uchar>(yc, xc);
    return static_cast<double>(1.0 * a + (x - xf) * (b - a));
  }

  template <typename T>
  double fabsN(const T& a, const T& b)
  {
    T diff = a - b;
    double d = 0.0;
    for (int i = 0; i < a.rows; ++i)
      d += fabs(a[i] - b[i]);
    return d / a.rows;
  }

  template <typename T>
  inline bool isBetween(const T &a, const T &L, const T &H) { return ((a >= L) && (a < H));};

  ros::NodeHandle nh_;
  ros::NodeHandle private_nh_;
  image_transport::ImageTransport it_;
  image_transport::SubscriberFilter imgL_sub_;
  image_transport::SubscriberFilter imgR_sub_;
  message_filters::Subscriber<sensor_msgs::CameraInfo> infoL_sub_;
  message_filters::Subscriber<sensor_msgs::CameraInfo> infoR_sub_;
  boost::shared_ptr<message_filters::Synchronizer<SyncPolicy> > sync_;
  ros::Subscriber trigger_sub_;
  ros::Subscriber smooth_sub_;
  ros::Subscriber scale_sub_;
  ros::Publisher disp_pub_;
  ros::Publisher cloud_pub_;

  boost::mutex mutex_;
  bool processing_;
  bool trigger_;

  image_geometry::StereoCameraModel model_;
  double scale_;
  double wsmooth_;
  int winsize_;
  int dispmax_, dispmin_;
  double distmax_, distmin_;
  bool do_sobel_;
};

} // namespace stereo

#endif  // _DENSE_STEREO_GRAPHCUT_H_
